{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optical Character Detection using TAO OCDNet\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\">\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Take a pretrained model and train an OCDNet model on ICDAR2015 dataset\n",
    "* Evaluate the trained model\n",
    "* Run inference with the trained model and visualize the result\n",
    "* Export the trained model to an .onnx file for deployment to DeepStream\n",
    "* Generate TensorRT engine using tao-deploy and verify the engine through evaluation\n",
    "\n",
    "At the end of this notebook, you will have generated a trained `ocdnet` model\n",
    "which you may deploy via [DeepStream](https://developer.nvidia.com/deepstream-sdk).\n",
    "\n",
    "For more information about OCDNet, you can take a look at [OCDNet](https://docs.nvidia.com/tao/tao-toolkit/text/object_detection/ocd.html) documentation page.\n",
    "\n",
    "Following is a sample prediction of OCDNet model.\n",
    "<!--- from img_5.jpg of ICDAR2015 test dataset -->\n",
    "<img align=\"center\" src=\"https://github.com/vpraveen-nv/model_card_images/blob/main/cv/notebook/ocdnet/img_5_result.jpg?raw=true\" width=\"640\">\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "This notebook shows an example usecase of OCDNet using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO launcher](#head-1)\n",
    "2. [Prepare dataset and pre-trained model](#head-2)\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Run TAO training](#head-4)\n",
    "5. [Evaluate trained models](#head-5)\n",
    "6. [Prune trained models](#head-6)\n",
    "7. [Retrain pruned models](#head-7)\n",
    "8. [Evaluate retrained models](#head-8)\n",
    "9. [Inferences](#head-9)\n",
    "10. [Deploy](#head-10)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "The following notebook requires the user to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the users workspace. Please note that the dataset to run this notebook is expected to reside in the `$LOCAL_PROJECT_DIR/data`, while the TAO experiment generated collaterals will be output to `$LOCAL_PROJECT_DIR/ocdnet/results`. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.\n",
    "\n",
    "The TAO launcher uses docker containers under the hood, and **for our data and results directory to be visible to the docker, they need to be mapped**. The launcher can be configured using the config file `~/.tao_mounts.json`. Apart from the mounts, you can also configure additional options like the Environment Variables and amount of Shared Memory available to the TAO launcher. <br>\n",
    "\n",
    "`IMPORTANT NOTE:` The code below creates a sample `~/.tao_mounts.json`  file. Here, we can map directories in which we save the data, specs, results and cache. You should configure it for your specific case so these directories are correctly visible to the docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=/path/to/local/tao-experiments\n",
    "\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\", \"ocdnet\")\n",
    "os.environ[\"HOST_RESULTS_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"ocdnet\", \"results\")\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/ocdnet\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"HOST_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir -p $HOST_DATA_DIR\n",
    "! mkdir -p $HOST_SPECS_DIR\n",
    "! mkdir -p $HOST_RESULTS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tlt_configs = {\n",
    "   \"Mounts\":[\n",
    "         # Mapping the Local project directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "        {\n",
    "           \"source\": os.environ[\"HOST_DATA_DIR\"],\n",
    "           \"destination\": \"/data/ocdnet\"\n",
    "        },\n",
    "        {\n",
    "           \"source\": os.environ[\"HOST_SPECS_DIR\"],\n",
    "           \"destination\": \"/specs\"\n",
    "        },\n",
    "        {\n",
    "           \"source\": os.environ[\"HOST_RESULTS_DIR\"],\n",
    "           \"destination\": \"/results\"\n",
    "        }\n",
    "   ],\n",
    "   \"DockerOptions\": {\n",
    "        \"shm_size\": \"16G\",\n",
    "        \"ulimits\": {\n",
    "            \"memlock\": -1,\n",
    "            \"stack\": 67108864\n",
    "         },\n",
    "        \"user\": \"{}:{}\".format(os.getuid(), os.getgid()),\n",
    "        \"network\": \"host\"\n",
    "   }\n",
    "}\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tlt_configs, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in the `nvidia-pyindex` python index. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 525.81+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python >=3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the virtualenv and virtualenvwrapper packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset and pre-trained models<a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Prepare dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " We will be using the `ICDAR2015` dataset for the ocdnet tutorial. Please access https://rrc.cvc.uab.es/?ch=4&com=tasks to register and download the data from `Task 4.1: Text Localization`. Unzip the files to `$HOST_DATA_DIR/`. \n",
    "\n",
    " The data will then be extracted to have below structure.\n",
    "\n",
    "```bash\n",
    "│── train\n",
    "│   ├──img\n",
    "|   ├──gt\n",
    "│── test\n",
    "│   ├──img\n",
    "|   ├──gt\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create local dir\n",
    "!mkdir -p $HOST_DATA_DIR/train/img\n",
    "!mkdir -p $HOST_DATA_DIR/train/gt\n",
    "!mkdir -p $HOST_DATA_DIR/test/img\n",
    "!mkdir -p $HOST_DATA_DIR/test/gt\n",
    "# unzip training data\n",
    "!unzip $HOST_DATA_DIR/ch4_training_images.zip -d $HOST_DATA_DIR/train/img\n",
    "!unzip $HOST_DATA_DIR/ch4_training_localization_transcription_gt.zip -d $HOST_DATA_DIR/train/gt\n",
    "# unzip test data\n",
    "!unzip $HOST_DATA_DIR/ch4_test_images.zip -d $HOST_DATA_DIR/test/img\n",
    "!unzip $HOST_DATA_DIR/Challenge4_Test_Task1_GT.zip -d $HOST_DATA_DIR/test/gt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verification. Training dataset contains 1000 images. Test dataset contains 500 images.\n",
    "!ls $HOST_DATA_DIR/train/img  |wc -l\n",
    "!ls $HOST_DATA_DIR/train/gt  |wc -l\n",
    "!ls $HOST_DATA_DIR/test/img  |wc -l\n",
    "!ls $HOST_DATA_DIR/test/gt  |wc -l"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__[Optional]__ If your image size is high resolution(such as 4000x4000) and you want to crop your images to small size for training, you can use following cell to run offline crop. This offline crop needs you to specify your excepted cropped image size(such as 800x800), it will generate cropped images and related labels to the `patch` folder under your dataset path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Offline crop for training dataset \n",
    "# !python3 offline_crop.py --dataset-path $HOST_DATA_DIR/train \\\n",
    "#                          --patch-height 640 \\\n",
    "#                          --patch-width 640 \\\n",
    "#                          --overlapRate 0.5 \\\n",
    "#                          --has-gt True \\\n",
    "#                          --img-ext jpg \\\n",
    "#                          --visible True"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Prepare pre-trained models\n",
    "\n",
    "We will use NGC CLI to get the pre-trained models. For more details, go to [ngc.nvidia.com](ngc.nvidia.com) and click the SETUP on the navigation bar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $HOST_RESULTS_DIR/ngccli\n",
    "\n",
    "# Remove any previously existing CLI installations\n",
    "!rm -rf $HOST_RESULTS_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $HOST_RESULTS_DIR/ngccli\n",
    "!unzip -u \"$HOST_RESULTS_DIR/ngccli/$CLI\" -d $HOST_RESULTS_DIR/ngccli/\n",
    "!rm $HOST_RESULTS_DIR/ngccli/*.zip \n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"HOST_RESULTS_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ngc registry model list nvidia/tao/ocdnet:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $HOST_RESULTS_DIR/pretrained_ocdnet/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pull pretrained model from NGC\n",
    "!ngc registry model download-version nvidia/tao/ocdnet:trainable_resnet18_v1.0 --dest $HOST_RESULTS_DIR/pretrained_ocdnet\n",
    "!ngc registry model download-version nvidia/tao/ocdnet:trainable_resnet50_v1.0 --dest $HOST_RESULTS_DIR/pretrained_ocdnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Check the models are downloaded into dir.\")\n",
    "!ls -l $HOST_RESULTS_DIR/pretrained_ocdnet/ocdnet_vtrainable_resnet18_v1.0\n",
    "!ls -l $HOST_RESULTS_DIR/pretrained_ocdnet/ocdnet_vtrainable_resnet50_v1.0"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-3\"></a>\n",
    "\n",
    "We provide specification files to configure the training parameters including:\n",
    "\n",
    "* train: configure the training hyperparameters\n",
    "    * num_gpus: number of gpus \n",
    "    * results_dir: Path to restore training result\n",
    "    * resume_training_checkpoint_path: Resume training from a checkpoint.\n",
    "    * num_epochs: The total epochs for training\n",
    "    * validation_interval: validation interval\n",
    "    * checkpoint_interval: checkpoint interval\n",
    "    * optimizer\n",
    "        * type: Defaults to Adam.\n",
    "        * lr: Initial learning rate \n",
    "    * lr_scheduler\n",
    "        * type: Only supports WarmupPolyLR\n",
    "        * warmup_epoch: The epoch numbers for warm up to initinal learning rate. It should be different from num_epochs. \n",
    "    * post_processing\n",
    "        * type: Only supports SegDetectorRepresenter\n",
    "        * thresh: The threshold for binarization.\n",
    "        * box_thresh: The threshold for bounding box.\n",
    "        * unclip_ratio: Default to 1.5. The box will look larger if this ratio is set to larger.\n",
    "    * Metric\n",
    "        * type: Only supports QuadMetric\n",
    "        * is_output_polygon: Defaults to false. False for bounding box. True for polygon.\n",
    "* dataset: configure the dataset and augmentation methods\n",
    "    * train_dataset:\n",
    "        * data_path: Path to train images. If there are multi sources, set it looks like ['/path/1' , '/path/2']\n",
    "        * pre_processes\n",
    "            * size: Ramdom crop size during training. Defaults to [640, 640].\n",
    "        * loader\n",
    "            * batch_size: batch size for dataloader\n",
    "            * num_workers: number of workers to do data loading \n",
    "    * validate_dataset: \n",
    "        * data_path: Path to validation images. If there are multi sources, set it looks like ['/path/1' , '/path/2']\n",
    "        * pre_processes\n",
    "            * short_size: Resize to width x height during evaluation. Defaults to [1280, 736].\n",
    "            * resize_text_polys: Resize the coordinate of text groudtruth. Defaults to true.\n",
    "        * loader\n",
    "            * batch_size: batch size for dataloader\n",
    "            * num_workers: number of workers to do data loading            \n",
    "* model: configure the model setting\n",
    "    * backbone: The backbone type. The deformable_resnet18 and deformable_resnet50 are supported.\n",
    "    * load_pruned_graph: Defaults to False. Must set to true if train a model which is pruned. \n",
    "    * pruned_graph_path: The path to the pruned model graph.\n",
    "    * pretrained_model_path: Finetune from a pretrained model. The `.pth` model is supported.\n",
    "\n",
    "Please refer to the TAO documentation about OCDNet to get more parameters that are configurable.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/train.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* Provide the sample spec file and the output directory location for models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: The following paths are set from the perspective of the TAO Docker.\n",
    "\n",
    "# The data is saved here\n",
    "%env DATA_DIR=/data/ocdnet\n",
    "%env SPECS_DIR=/specs\n",
    "%env RESULTS_DIR=/results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!echo $HOST_DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#print(\"Run training with ngc pretrained model. \")\n",
    "#Please note that this training will take a while if using default specs (30 epochs).\n",
    "\n",
    "!tao model ocdnet train \\\n",
    "          -e $SPECS_DIR/train.yaml \\\n",
    "          results_dir=$RESULTS_DIR \\\n",
    "          model.pretrained_model_path=$RESULTS_DIR/pretrained_ocdnet/ocdnet_vtrainable_resnet18_v1.0/ocdnet_deformable_resnet18.pth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#print(\"Resume training with the checkpoint corresponding to your set epoch.\")\n",
    "#For example, resume training from 3rd checkpoint as below.\n",
    "\n",
    "#%env NUM_EPOCH=003\n",
    "#!cp \"`ls -rlt $HOST_RESULTS_DIR/train/ocd_model_epoch\\=${NUM_EPOCH}*.pth |tail -1 |awk -F \" \" '{print $NF}'`\" $HOST_RESULTS_DIR/train/resume.pth\n",
    "\n",
    "#!tao model ocdnet train \\\n",
    "#          -e $SPECS_DIR/train.yaml \\\n",
    "#          results_dir=$RESULTS_DIR \\\n",
    "#          train.resume_training_checkpoint_path=$RESULTS_DIR/train/resume.pth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"For multi-GPU, change train.num_gpus in train.yaml or set train.num_gpus in commandline based on your machine.\")\n",
    "#For example, run training with 2gpus.\n",
    "#!tao model ocdnet train \\\n",
    "#          -e $SPECS_DIR/train.yaml \\\n",
    "#          results_dir=$RESULTS_DIR \\\n",
    "#          model.pretrained_model_path=$RESULTS_DIR/pretrained_ocdnet/ocdnet_vtrainable_resnet18_v1.0/ocdnet_deformable_resnet18.pth \\\n",
    "#          train.num_gpus=2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Trained checkpoints:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "# %env NUM_EPOCH=029\n",
    "\n",
    "# Get the name of the checkpoint corresponding to your set epoch\n",
    "# tmp=!ls $HOST_RESULTS_DIR/train/*.pth | grep epoch_$NUM_EPOCH\n",
    "# %env CHECKPOINT={tmp[0]}\n",
    "\n",
    "# Or get the latest checkpoint\n",
    "os.environ[\"CHECKPOINT\"] = os.path.join(os.getenv(\"HOST_RESULTS_DIR\"), \"train/ocd_model_latest.pth\")\n",
    "\n",
    "print('Rename a trained model: ')\n",
    "print('---------------------')\n",
    "!cp $CHECKPOINT $HOST_RESULTS_DIR/train/ocd_model.pth\n",
    "!ls -ltrh $HOST_RESULTS_DIR/train/ocd_model.pth"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate a trained model <a class=\"anchor\" id=\"head-5\"></a>\n",
    "\n",
    "In this section, we run the `evaluate` tool to evaluate the trained model and produce the evaluation metric. The evaluation metric will show the hmeans of different threshold for binarization, you could select the threshold which has the best hmean result.\n",
    "\n",
    "We provide specification files to configure the parameters including:\n",
    "\n",
    "* evaluate: configure the training hyperparameters\n",
    "    * results_dir: Path to restore training result\n",
    "    * checkpoint: checkpoint path for running evaluation\n",
    "    * post_processing\n",
    "        * type: Only supports SegDetectorRepresenter\n",
    "        * box_thresh: The threshold for bounding box.\n",
    "        * unclip_ratio: Default to 1.5. The box will look larger if this ratio is set to larger.\n",
    "    * Metric\n",
    "        * type: Only supports QuadMetric\n",
    "        * is_output_polygon: Defaults to false. False for bounding box. True for polygon.\n",
    "* dataset: configure the dataset and augmentation methods\n",
    "    * validate_dataset: \n",
    "        * data_path: Path to validation images. If there are multi sources, set it looks like ['/path/1' , '/path/2']\n",
    "        * pre_processes\n",
    "            * short_size: Resize to width x height during evaluation. Defaults to [1280, 736].\n",
    "            * resize_text_polys: Resize the coordinate of text groudtruth. Defaults to true.\n",
    "        * loader\n",
    "            * batch_size: batch size for dataloader\n",
    "            * num_workers: number of workers to do data loading            \n",
    "* model: configure the model setting\n",
    "    * backbone: The backbone type. The deformable_resnet18 and deformable_resnet50 are supported.\n",
    "    * load_pruned_graph: Whether evaluation a model which has pruned model graph. Defaults to False.\n",
    "    * pruned_graph_path: The path to the pruned model graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate on model\n",
    "!tao model ocdnet evaluate \\\n",
    "            -e $SPECS_DIR/evaluate.yaml \\\n",
    "            evaluate.checkpoint=$RESULTS_DIR/train/ocd_model.pth"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 6. Prune trained models <a class=\"anchor\" id=\"head-6\"></a>\n",
    "\n",
    "In this section, we run the `prune` tool to get a pruned model.\n",
    "\n",
    "We provide specification files to configure as following:\n",
    "\n",
    "* prune: configure the pruning hyperparameters\n",
    "    * results_dir: Path to restore training result\n",
    "    * checkpoint: Checkpoint path for running pruning\n",
    "    * ch_sparsity: Channel sparisty. Also known as pruning threshold. Higher value gives you smaller model.\n",
    "    * p: The norm degree, Default: 2. By default, it calculates the group L2-norm for each channel/dim.\n",
    "    * round_to: Round channels to the nearest multiple of round_to. E.g., round_to=8 means channels will be rounded to 8x. Default: 32.\n",
    "* dataset: configure the dataset and augmentation methods\n",
    "    * validate_dataset: \n",
    "        * data_path: Path to validation images. If there are multi sources, set it looks like ['/path/1' , '/path/2']\n",
    "        * pre_processes\n",
    "            * short_size: Resize to width x height during evaluation. Defaults to [1280, 736].\n",
    "            * resize_text_polys: Resize the coordinate of text groudtruth. Defaults to true.\n",
    "        * loader\n",
    "            * batch_size: batch size for dataloader\n",
    "            * num_workers: number of workers to do data loading            \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%env ch_sparsity=0.1\n",
    "\n",
    "!tao model ocdnet prune -e $SPECS_DIR/prune.yaml \\\n",
    "                prune.checkpoint=$RESULTS_DIR/train/ocd_model.pth \\\n",
    "                prune.ch_sparsity=$ch_sparsity \\\n",
    "                prune.results_dir=$RESULTS_DIR/prune"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -lht $HOST_RESULTS_DIR/prune/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " ## 7. Retrain pruned models <a class=\"anchor\" id=\"head-7\"></a>\n",
    " * Model needs to be re-trained to bring back accuracy after pruning\n",
    " * Need to set load_pruned_graph to __True__ and set the path of the pruned graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retraining using the pruned graph \n",
    "!tao model ocdnet train -e $SPECS_DIR/train.yaml \\\n",
    "                  train.results_dir=$RESULTS_DIR/retrain \\\n",
    "                  model.load_pruned_graph=true \\\n",
    "                  model.pruned_graph_path=$RESULTS_DIR/prune/pruned_$ch_sparsity.pth\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Listing the newly retrained model.\n",
    "!ls -lht $HOST_RESULTS_DIR/retrain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "# %env NUM_EPOCH=029\n",
    "\n",
    "# Get the name of the checkpoint corresponding to your set epoch\n",
    "# tmp=!ls $HOST_RESULTS_DIR/retrain/*.pth | grep epoch_$NUM_EPOCH\n",
    "# %env CHECKPOINT={tmp[0]}\n",
    "\n",
    "# Or get the latest checkpoint\n",
    "os.environ[\"CHECKPOINT\"] = os.path.join(os.getenv(\"HOST_RESULTS_DIR\"), \"retrain/ocd_model_latest.pth\")\n",
    "\n",
    "print('Rename a retrained model: ')\n",
    "print('---------------------')\n",
    "!cp $CHECKPOINT $HOST_RESULTS_DIR/retrain/ocd_model.pth\n",
    "!ls -ltrh $HOST_RESULTS_DIR/retrain/ocd_model.pth"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Evaluate retrained model <a class=\"anchor\" id=\"head-8\"></a>\n",
    "* Need to set load_pruned_graph to true and set the path of the pruned graph.\n",
    "* Model pruning reduces model parameters to improve inference frames per second (FPS) while maintaining nearly the same hmean."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model ocdnet evaluate -e $SPECS_DIR/evaluate.yaml \\\n",
    "                     evaluate.checkpoint=$RESULTS_DIR/retrain/ocd_model.pth \\\n",
    "                     model.load_pruned_graph=true \\\n",
    "                     model.pruned_graph_path=$RESULTS_DIR/prune/pruned_$ch_sparsity.pth\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# According the evaluation results, set threshold which has the best hmean result.\n",
    "%env best_thresh=FIXME"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Visualize Inferences <a class=\"anchor\" id=\"head-9\"></a>\n",
    "In this section, we run the `inference` tool to generate inferences on the trained models and visualize the results. The `inference` tool produces annotated image outputs and txt files that contain prediction information.\n",
    "\n",
    "We provide specification files to configure the inference parameters including:\n",
    "\n",
    "* inference: configure the training hyperparameters\n",
    "    * results_dir: Path to restore inference result\n",
    "    * checkpoint: checkpoint path for running inference\n",
    "    * input_folder: The input folder for inference\n",
    "    * width: The width for resizing\n",
    "    * height: The height for resizing\n",
    "    * polygon: Produce polygon(true) or bounding box(false). Defaults to false.\n",
    "    * post_processing\n",
    "        * type: Only supports SegDetectorRepresenter\n",
    "        * thresh: The threshold for binarization.\n",
    "        * box_thresh: The threshold for bounding box.\n",
    "        * unclip_ratio: Default to 1.5. The box will look larger if this ratio is set to larger.       \n",
    "* model: configure the model setting\n",
    "    * backbone: The backbone type. The deformable_resnet18 and deformable_resnet50 are supported.\n",
    "    * load_pruned_graph: Whether evaluation a model which has pruned model graph. Defaults to False.\n",
    "    * pruned_graph_path: The path to the pruned model graph.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model ocdnet inference \\\n",
    "        -e $SPECS_DIR/inference.yaml \\\n",
    "        inference.checkpoint=$RESULTS_DIR/train/ocd_model.pth \\\n",
    "        inference.input_folder=$DATA_DIR/test/img \\\n",
    "        inference.results_dir=$RESULTS_DIR/infer/ \\\n",
    "        inference.post_processing.args.thresh=$best_thresh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple grid visualizer\n",
    "!pip3 install \"matplotlib>=3.3.3, <4.0\"\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from math import ceil\n",
    "result_image = ['result.jpg']\n",
    "\n",
    "def visualize_images(output_path, num_cols=4, num_images=10):\n",
    "    num_rows = int(ceil(float(num_images) / float(num_cols)))\n",
    "    f, axarr = plt.subplots(num_rows, num_cols, figsize=[80,30])\n",
    "    f.tight_layout()\n",
    "    a = [os.path.join(output_path, image) for image in sorted(os.listdir(output_path)) \n",
    "         if image.split(\"_\")[-1] in result_image]\n",
    "    for idx, img_path in enumerate(a[:num_images]):\n",
    "        col_id = idx % num_cols\n",
    "        row_id = idx // num_cols\n",
    "        img = plt.imread(img_path)\n",
    "        axarr[row_id, col_id].imshow(img) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    " # Visualizing the sample images.\n",
    "IMAGE_DIR = os.path.join(os.environ['HOST_RESULTS_DIR'], \"infer\")\n",
    "COLS = 2 # number of columns in the visualizer grid.\n",
    "IMAGES = 4 # number of images to visualize.\n",
    "\n",
    "visualize_images(IMAGE_DIR, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10. Deploy <a class=\"anchor\" id=\"head-10\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Export pth model to ONNX model\n",
    "!tao model ocdnet export \\\n",
    "           -e $SPECS_DIR/export.yaml \\\n",
    "           export.checkpoint=$RESULTS_DIR/train/ocd_model.pth \\\n",
    "           export.onnx_file=$RESULTS_DIR/export/ocd_model.onnx\n",
    "           "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate TensorRT engine using tao-deploy\n",
    "!tao deploy ocdnet gen_trt_engine -e $SPECS_DIR/gen_trt_engine.yaml \\\n",
    "                               gen_trt_engine.onnx_file=$RESULTS_DIR/export/ocd_model.onnx \\\n",
    "                               gen_trt_engine.tensorrt.min_batch_size=1 \\\n",
    "                               gen_trt_engine.tensorrt.opt_batch_size=3 \\\n",
    "                               gen_trt_engine.tensorrt.max_batch_size=3 \\\n",
    "                               gen_trt_engine.tensorrt.data_type=int8 \\\n",
    "                               gen_trt_engine.trt_engine=$RESULTS_DIR/export/ocdnet_model.engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate with generated TensorRT engine\n",
    "%env CUDA_MODULE_LOADING=\"LAZY\"\n",
    "!tao deploy ocdnet evaluate -e $SPECS_DIR/evaluate.yaml \\\n",
    "                             evaluate.trt_engine=$RESULTS_DIR/export/ocdnet_model.engine \\\n",
    "                             evaluate.post_processing.args.thresh=$best_thresh\n",
    "                             "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inference with generated TensorRT engine\n",
    "!tao deploy ocdnet inference -e $SPECS_DIR/inference.yaml \\\n",
    "                              inference.trt_engine=$RESULTS_DIR/export/ocdnet_model.engine \\\n",
    "                              inference.input_folder=$DATA_DIR/test/img \\\n",
    "                              inference.results_dir=$RESULTS_DIR/inference \\\n",
    "                              inference.post_processing.args.thresh=$best_thresh"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook has come to an end."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
